import datetime
import pytz

from bson import ObjectId, code
from flask import g
from wtforms import fields, ValidationError
from wtforms.validators import DataRequired

from app.config.settings import DATETIME_FMT
from app.db.models.assets import Asset
from app.db.models.jobs import Job
from app.db.models.tasks import Task
from app.errors import JobNotFound, TaskNotFound, ParameterError
from app.validators import BaseForm
from app.libs.enums import WS_TASK_GROUP_TUPLE, TaskType


class ReportDownloadForm(BaseForm):
    job_id = fields.StringField(validators=(DataRequired(),))
    report_type = fields.StringField(validators=(DataRequired(),))
    start_date = fields.DateTimeField()
    end_date = fields.DateTimeField()

    def validate_job_id(self, value):
        if job_id := value.data:
            try:
                value.data = ObjectId(job_id)
            except Exception:
                raise ValidationError('job_id format error')
            else:
                self.job = Job.objects.filter(uid=g.user.id, id=value.data).first()
                # self.job = Job.objects.filter(id=value.data).first()
                if not self.job:
                    raise JobNotFound()

    def validate_report_type(self, value):
        if value.data not in ["latest", "previous", "range"]:
            raise ValidationError("下载范围参数错误")

    def get_task_type_mapping(self) -> dict:
        """
        获取三类["latest", "previous", "range"]，对应最新监测、上一次监测、时间范围监测三种筛选条件的任务id

        Returns: 任务类对应的任务数据:
        {"ssl": [mongoObject], "vul": [], "securityEvent": [], "http": [], "ping": [], "content": [], "asset": []}

        """
        params = self.validate_for_api()
        report_type = params.data["report_type"]
        job_id = params.data["job_id"]

        normal_task_type_tuple = ('ssl', 'vul', 'securityEvent', 'http', 'ping', 'content')
        special_task_type_tuple = ("asset",)

        project = {
            "taskSettings": 0,
            "alertSettings": 0,
            "triggerType": 0,
            "isSpecialTask": 0,
            "progress": 0,
            "isLastTask": 0,

            "taskIdx": 0,
            "refType": 0,
            "refId": 0,
            "taskId": 0,
            "nextExecDate": 0,
            "result.addition.detail": 0
        }
        # 最近监测的一般字段为【任务类型】+ResultId
        if report_type in ["latest", "previous"]:
            self.task_type_ids = []
            # 获取除asset之外监测类型的其他监测结果
            for task_type in normal_task_type_tuple:
                task_id_field = f"{task_type}ResultId" if report_type == "latest" else f"{task_type}PrevResultId"
                task_id = self.job.to_mongo().get(task_id_field)
                if isinstance(task_id, ObjectId):
                    self.task_type_ids.append(task_id)
                elif isinstance(task_id, list):
                    if task_id and task_id[0]:
                        self.task_type_ids.extend(task_id)
            # 单独获取asset的上一次任务id,逻辑与其他类型不一致
            asset_id = self.job.assetId
            asset = Asset.objects.filter(id=asset_id).first()
            asset_task_id_field = "resultId" if report_type == "latest" else "PrevResultId"
            asset_task_id = asset.to_mongo().get(asset_task_id_field)
            if asset_task_id:
                self.task_type_ids.append(asset_task_id)

            tasks = Task._get_collection().find({"_id": {"$in": self.task_type_ids}, 'status': 'completed'}, project)
        # 如果是按照时间范围筛选，则直接查询监测结束时间的时间范围任务
        else:
            start_date = datetime.datetime.strptime(params.data["start_date"], DATETIME_FMT)
            end_date = datetime.datetime.strptime(params.data["end_date"], DATETIME_FMT)
            self.query_start_date = start_date.astimezone(pytz.timezone('Asia/Shanghai'))
            self.query_end_date = end_date.astimezone(pytz.timezone('Asia/Shanghai'))

            query = {
                "jobId": job_id,
                "status": "completed",
                "startTime": {
                    "$gte": self.query_start_date,
                },
                "endTime": {
                    "$lte": self.query_end_date
                }
            }
            tasks = Task._get_collection().find(query, project)

        all_task_type = normal_task_type_tuple + special_task_type_tuple
        task_type_mapping = {_type: [] for _type in all_task_type}
        for task in tasks:
            if task.get('taskType') not in all_task_type:  # fix websoc_6851 问题
                continue
            task_type_mapping[task.get("taskType")].append(task)
        return task_type_mapping

    def get_http_node_data(self):
        http_error_node_continuance = self._aggregate_http_ping_node_error_continuance(task_type="http")
        http_node_detail = self._aggregate_error_node_detail(task_type="http")
        return {
            "node_continuance": http_error_node_continuance,
            "node_detail": http_node_detail
        }

    def get_ping_node_data(self):
        ping_error_node_continuance = self._aggregate_http_ping_node_error_continuance(task_type="ping")
        ping_node_detail = self._aggregate_error_node_detail(task_type="ping")
        return {
            "node_continuance": ping_error_node_continuance,
            "node_detail": ping_node_detail
        }

    def _aggregate_http_ping_node_error_continuance(self, task_type=None):

        # 在mongodb服务器中执行的$function js语句，用于计算数组中异常节点持续时间
        js_code = code.Code(
            """
      function get_node_error_time (node_time_list) {
          const error_list = []
          let start_index = -1
          let end_index = -1
          let end_time = null
          let start_time = null
          const final_elem_index = node_time_list.length - 1
          node_time_list.sort((a, b) => {
            return new Date(b.monitor_time) - new Date(a.monitor_time)
          }).reverse().map((item, index) => {
            const has_error = item.securityStatus === "warning"

            if (has_error && start_index === -1) {

              start_index = index
              start_time = item.monitor_time
            }
            if (!has_error && start_index !== -1) {
              end_index = index
            } else if (has_error && start_index !== -1 && index === final_elem_index) {
              end_index = index
            }
            if (end_index !== -1) {
              end_time = item.monitor_time
            }

            if (start_time && end_time) {

              error_item = {
                area: item.area,
                ip: item.ip,
                avgTime: (node_time_list.slice(start_index, end_index).reduce((total, i) => {
                  total += Number(i.totalTime)
                  return total
                }, 0) / (end_index - start_index)).toFixed(2),
                continuanceTime: (end_time - start_time) / 1000,
                timeRange: { start_time, end_time }
              }
              start_index = -1
              end_index = -1
              start_time = null
              end_time = null
              error_list.push(error_item)
            }
          })
          return error_list
        }
            """
        )
        if hasattr(self, "task_type_ids"):
            match = {
                "$match": {
                    "_id": {"$in": self.task_type_ids}
                }
            }
        else:
            match = {
                "$match": {
                    "jobId": self.job.id,
                    "taskType": task_type,
                    "startTime": {
                        "$gte": self.query_start_date,
                        "$lte": self.query_end_date
                    }
                }
            }

        pipeline = [
            match,
            {
                "$project": {
                    "result.addition.detail.area": 1,
                    "result.addition.detail.securityStatus": 1,
                    "result.addition.detail.ip": 1,
                    "result.addition.detail.totalTime": 1,
                    "result.addition.detail.monitor_time": "$startTime",
                    "result.addition.detail.target": "$target"

                }
            },
            {
                "$unwind": "$result.addition.detail"
            },
            {
                "$group": {
                    # "_id": "$result.addition.detail.area",
                    "_id": {"area": "$result.addition.detail.area", "target": "$result.addition.detail.target"},
                    "additions": {
                        "$push": "$result.addition.detail"
                    }
                }
            },
            {
                "$project": {
                    "result": {
                        "$function": {
                            "body": code.Code(js_code),
                            "args": ["$additions"],
                            "lang": "js"
                        }
                    }
                }
            }
        ]

        # result = Task.objects.aggregate(pipeline)
        result = Task._get_collection().aggregate(pipeline)
        return list(result)

    def _aggregate_error_node_detail(self, task_type=None):
        if hasattr(self, "task_type_ids"):
            match = {
                "$match": {
                    "_id": {"$in": self.task_type_ids}
                }
            }
        else:
            match = {
                "$match": {
                    "jobId": self.job.id,
                    "taskType": task_type,
                    "startTime": {
                        "$gte": self.query_start_date,
                        "$lte": self.query_end_date
                    }
                }
            }
        pipeline = [
            match,
            {
                "$project": {
                    "result.addition.detail.area": 1,
                    "result.addition.detail.securityStatus": 1,
                    "result.addition.detail.ip": 1,
                    "result.addition.detail.totalTime": 1,
                    "result.addition.detail.has_error": 1,
                    "result.addition.detail.monitor_time": "$startTime",
                    "result.addition.detail.target": "$target",
                }
            },
            {
                "$unwind": "$result.addition.detail"
            },
            {
                "$match": {
                    "result.addition.detail.securityStatus": "warning"
                }
            },
            {
                "$group": {
                    "_id": {
                        "target": "$result.addition.detail.target",
                        "area": "$result.addition.detail.area"
                    },
                    "additions": {
                        "$push": "$result.addition.detail"
                    }
                }
            }
        ]
        result = Task._get_collection().aggregate(pipeline)
        return result


class ReportDownloadTaskForm(ReportDownloadForm):
    uid = fields.StringField()
    job_id = fields.StringField()
    job_ids = fields.StringField(validators=(DataRequired(),))
    download_type = fields.StringField(validators=(DataRequired(),))
    file_type = fields.StringField()  # 下载文件类型，docx, zip

    def validate_uid(self, value):
        if g.role.name == '管理员':
            try:
                self.uid = ObjectId(value.data)
            except Exception:
                raise ParameterError(msg='用户信息错误')

    def validate_job_ids(self, value):
        if job_ids_str := value.data:
            try:
                value.data = [ObjectId(job_id) for job_id in job_ids_str.split(',')]
            except Exception:
                raise ValidationError('job_ids format error')
            else:
                uid = g.user.id
                if g.role.name == '管理员':
                    uid = self.uid
                self.jobs = Job.objects.filter(uid=uid, id__in=value.data)
                if self.jobs.count() == 0:
                    raise JobNotFound()

    def validate_download_type(self, value):
        if value.data not in ["0", "1"]:
            raise ValidationError("下载类型参数错误")


class OneTaskReportDownloadForm(BaseForm):
    task_id = fields.StringField(validators=(DataRequired(),))

    def validate_task_id(self, value):
        project = {
            "taskSettings": 0,
            "alertSettings": 0,
            "triggerType": 0,
            "isSpecialTask": 0,
            "progress": 0,
            "isLastTask": 0,
            "taskIdx": 0,
            "refType": 0,
            "refId": 0,
            "taskId": 0,
            "nextExecDate": 0,
            "result.addition.detail": 0
        }
        all_task_type = WS_TASK_GROUP_TUPLE + (TaskType.ipv6.value, TaskType.asset.value)
        self.task_type_mapping = {_type: [] for _type in all_task_type}
        if task_id := value.data:
            try:
                value.data = ObjectId(task_id)
            except Exception:
                raise ValidationError('task_id format error')
            else:
                self.task = Task._get_collection().find_one(
                    {"_id": value.data, "uid": g.user.id, "taskType": {"$in": all_task_type}},
                    project)
                if not self.task:
                    raise TaskNotFound()
                if self.task.get("status") != "completed":
                    raise ValidationError('task not completed')
                self.task_type_mapping[self.task["taskType"]].append(self.task)
                self.job = Job.objects.filter(id=self.task["jobId"]).first()


class ExportListForm(BaseForm):
    uid = fields.StringField()
    username = fields.StringField()
    page = fields.IntegerField()
    limit = fields.IntegerField()
    asset = fields.StringField()
    status = fields.StringField()
    source_ip = fields.StringField()

    def validate_uid(self, value):
        if g.role.name == '管理员':
            try:
                self.uid = ObjectId(value.data)
            except Exception:
                raise ParameterError(msg='用户信息错误')

    def validate_status(self, value):
        if value.data and value.data not in ['normal', 'pending', 'error', 'success']:
            raise ValidationError(msg='status value error')


class ReportDeleteForm(BaseForm):
    uid = fields.StringField()
    download_task_ids = fields.StringField()

    def validate_uid(self, value):
        if g.role.name == '管理员':
            try:
                self.uid = ObjectId(value.data)
            except Exception:
                raise ParameterError(msg='用户信息错误')

    def validate_download_task_ids(self, value):
        if value.data:
            value.data = [ObjectId(d) for d in value.data.split(",")]
